from crewai import Agent, Crew, Process, Task
from crewai.project import CrewBase, agent, crew, task
from product_recommendation.tools.database_tool import DatabaseTool
from product_recommendation.tools.nlp_tool import NLPTool
from product_recommendation.tools.base_tool import create_tool
from langchain_openai import ChatOpenAI

# 取消注释以下行以使用自定义工具的示例
# 从 product_recommendation.tools.custom_tool 导入 MyCustomTool

# 查看我们的工具文档以获取有关如何使用它们的更多信息
# from crewai_tools import SerperDevTool

@CrewBase
class ProductRecommendationCrew():
	"""商品推荐助手团队"""
	agents_config = 'config/agents.yaml'
	tasks_config = 'config/tasks.yaml'

	def __init__(self, llm: ChatOpenAI = None):
		self.db_tool = DatabaseTool()
		self.nlp_tool = NLPTool(llm)
		self.llm = llm

	@agent
	def intent_analyzer(self) -> Agent:
		return Agent(
			config=self.agents_config['intent_analyzer'],
			tools=[
				create_tool(
					name="analyze_intent",
					description="分析用户查询意图，返回结构化的意图信息",
					func=self.nlp_tool.analyze_intent
				)
			],
			llm=self.llm,
			verbose=True
		)

	@agent
	def product_researcher(self) -> Agent:
		return Agent(
			config=self.agents_config['product_researcher'],
			tools=[
				create_tool(
					name="get_products_by_category",
					description="根据类别获取商品列表",
					func=self.db_tool.get_products_by_category
				),
				create_tool(
					name="get_products_by_price_range",
					description="根据价格范围获取商品列表",
					func=self.db_tool.get_products_by_price_range
				),
				create_tool(
					name="search_products",
					description="根据关键词搜索商品",
					func=self.db_tool.search_products
				),
				create_tool(
					name="get_popular_products",
					description="获取热门商品列表",
					func=self.db_tool.get_popular_products
				)
			],
			llm=self.llm,
			verbose=True
		)

	@agent
	def recommendation_writer(self) -> Agent:
		return Agent(
			config=self.agents_config['recommendation_writer'],
			tools=[
				create_tool(
					name="generate_recommendation_text",
					description="根据商品信息生成推荐文案",
					func=self.nlp_tool.generate_recommendation_text
				)
			],
			llm=self.llm,
			verbose=True
		)

	@task
	def analyze_intent_task(self) -> Task:
		return Task(
			description="分析用户查询意图",
			agent=self.intent_analyzer,
			expected_output="查询意图分析结果",
			context=lambda inputs: f"用户查询: {inputs['query']}"
		)

	@task
	def research_products_task(self) -> Task:
		return Task(
			description="根据意图查找商品",
			agent=self.product_researcher,
			expected_output="商品列表",
			context=lambda inputs: f"根据意图 {inputs.get('intent', {})} 查找相关商品"
		)

	@task
	def generate_recommendation_task(self) -> Task:
		return Task(
			description="生成推荐文案",
			agent=self.recommendation_writer,
			expected_output="推荐文案",
			context=lambda inputs: f"为查询 '{inputs['query']}' 生成推荐文案"
		)

	@crew
	def crew(self) -> Crew:
		"""创建商品推荐助手团队"""
		return Crew(
			agents=self.agents,
			tasks=self.tasks,
			process=Process.sequential,
			verbose=True,
		)